package com.wang.judge.config;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.FileUtil;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.*;
import com.github.dockerjava.api.model.*;
import com.wang.common.core.constants.JudgeConstants;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

/**
 * 实现容器池,避免因为创建容器而产生的开销
 */
@Slf4j
public class DockerSandBoxPool {

    private DockerClient dockerClient; //容器客户端

    private String sandboxImage; //代码沙箱镜像

    private String volumeDir; //挂载目录,与宿主机中的目录进行关联

    private Long memoryLimit; //最大内存限制

    private Long memorySwapLimit; //最大内存交换次数限制

    private Long cpuLimit; //使用的最大的cpu核心数

    private int poolSize; //容器池中的最大容器数量

    private String containerNamePrefix; //容器前缀名

    private BlockingQueue<String> containerQueue; //归还与获取容器的阻塞队列

    private Map<String, String> containerNameMap;

    public DockerSandBoxPool(DockerClient dockerClient,
                             String sandboxImage,
                             String volumeDir, Long memoryLimit,
                             Long memorySwapLimit, Long cpuLimit,
                             int poolSize, String containerNamePrefix) {
        this.dockerClient = dockerClient;
        this.sandboxImage = sandboxImage;
        this.volumeDir = volumeDir;
        this.memoryLimit = memoryLimit;
        this.memorySwapLimit = memorySwapLimit;
        this.cpuLimit = cpuLimit;
        this.poolSize = poolSize;
        this.containerQueue = new ArrayBlockingQueue<>(poolSize);
        this.containerNamePrefix = containerNamePrefix;
        this.containerNameMap = new HashMap<>();
    }

    /**
     * 为容器池中创建容器
     */
    public void initDockerPool() {
        log.info("------  创建容器开始  -----");
        for(int i = 0; i < poolSize; i++) {
            createContainer(containerNamePrefix + "-" + i);
        }
        log.info("------  创建容器结束  -----");
    }

    /**
     * 从容器池中获取容器
     * @return 返回获取到的容器
     */
    public String getContainer() {
        try {
            return containerQueue.take();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 使用容器之后把容器归还
     * @param containerId 容器id
     */
    public void returnContainer(String containerId) {
        containerQueue.add(containerId);
    }

    /**
     * 为容器池中创建容器
     * @param containerName 容器名称
     */
    private void createContainer(String containerName) {
        List<Container> containerList = dockerClient.listContainersCmd().withShowAll(true).exec();
        // 检查容器列表是否非空
        if (!CollectionUtil.isEmpty(containerList)) {
            // 拼接完整的容器名称：前缀 + 容器名
            String names = JudgeConstants.JAVA_CONTAINER_PREFIX + containerName;
            // 遍历所有容器
            for (Container container : containerList) {
                // 获取容器的所有名称（Docker容器可以有多个别名）
                String[] containerNames = container.getNames();
                // 检查容器是否有名称，并且第一个名称是否与我们要创建的容器名称匹配
                if (containerNames != null && containerNames.length > 0 && names.equals(containerNames[0])) {
                    // 如果容器处于"created"(已创建但未启动)或"exited"(已停止)状态
                    if ("created".equals(container.getState()) || "exited".equals(container.getState())) {
                        // 启动这个容器
                        dockerClient.startContainerCmd(container.getId()).exec();
                    }
                    // 将容器ID加入容器池队列
                    containerQueue.add(container.getId());
                    // 在映射表中记录容器ID和名称的对应关系
                    containerNameMap.put(container.getId(), containerName);
                    // 直接返回，不再创建新容器
                    return;
                }
            }
        }
        //拉取镜像
        pullJavaEnvImage();
        //创建容器  限制资源  控制权限
        HostConfig hostConfig = getHostConfig(containerName);
        //创建容器
        CreateContainerCmd containerCmd = dockerClient
                .createContainerCmd(JudgeConstants.JAVA_ENV_IMAGE)
                .withName(containerName);
        //配置当前容器
        CreateContainerResponse createContainerResponse = containerCmd
                .withHostConfig(hostConfig)
                .withAttachStderr(true)
                .withAttachStdout(true)
                .withTty(true)
                .exec();
        //记录容器id
        String containerId = createContainerResponse.getId();
        //启动容器
        dockerClient.startContainerCmd(containerId).exec();
        containerQueue.add(containerId);
        containerNameMap.put(containerId, containerName);
    }

    /**
     * 创建镜像
     */
    private void pullJavaEnvImage() {
        //获取所有的镜像
        ListImagesCmd listImagesCmd = dockerClient.listImagesCmd();
        List<Image> imageList = listImagesCmd.exec();
        //判断要创建的镜像在已经创建的镜像中是否存在
        for (Image image : imageList) {
            String[] repoTags = image.getRepoTags();
            if (repoTags != null && repoTags.length > 0 && sandboxImage.equals(repoTags[0])) {
                return;//如果存在直接返回
            }
        }
        //如果不存在,就创建镜像
        PullImageCmd pullImageCmd = dockerClient.pullImageCmd(sandboxImage);
        try {
            //使用回调等待拉取完成,如果其中出现了异常,则抛出异常
            pullImageCmd.exec(new PullImageResultCallback()).awaitCompletion();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 限制容器调用操作系统的资源,限制容器权限
     * @return 返回配置资源和权限之后的容器配置
     */
    private HostConfig getHostConfig(String containerName) {
        HostConfig hostConfig = new HostConfig();
        //设置挂载目录，指定用户代码路径,与本地路径进行映射
        String userCodeDir = createContainerDir(containerName);
        hostConfig.setBinds(new Bind(userCodeDir, new Volume(volumeDir)));
        //限制docker容器使用资源
        hostConfig.withMemory(memoryLimit);//限制内存使用
        hostConfig.withMemorySwap(memorySwapLimit);//限制内存交换次数
        hostConfig.withCpuCount(cpuLimit);//限制使用cpu核心数
        hostConfig.withNetworkMode("none");  //禁用网络
        hostConfig.withReadonlyRootfs(true); //禁止在root目录写文件
        return hostConfig;
    }

    /**
     * 获取容器的挂载目录
     * @param containerId 容器id
     * @return 返回容器的挂载目录
     */
    public String getCodeDir(String containerId) {

        String containerName = containerNameMap.get(containerId); //获取容器名称

        log.info("containerName：{}", containerName); //日志打印容器名称

        return System.getProperty("user.dir") + File.separator + JudgeConstants.CODE_DIR_POOL + File.separator + containerName;
    }

    /**
     * 为每个容器指定挂载目录,与本地目录进行映射
     * @param containerName 容器名
     * @return 返回创建好的挂载目录
     */
    private String createContainerDir(String containerName) {
        //一级目录,存放所有容器的挂载目录
        String codeDir = System.getProperty("user.dir") + File.separator + JudgeConstants.CODE_DIR_POOL;
        //判断挂载目录是否存在,如果不存在,创建挂在目录,如果已经存在,则不创建
        if (!FileUtil.exist(codeDir)) {
            FileUtil.mkdir(codeDir);
        }
        return codeDir + File.separator + containerName;
    }
}
